#include "IndexHNSW.h"
#include <cassert>
#include <cmath>
#include <cstddef>
#include <utility>

IndexHNSW::IndexHNSW(size_t d, 
                     size_t n, 
                     std::unique_ptr<float[]> embedding,
                     int m,
                     int ef_construction,
                     int seed)
    : d_(d), ntotal_(n), embedding_(std::move(embedding)),
      m_(m), m0_(2 * m), ef_construction_(ef_construction), mt_(seed) {
        set_default_prob(1.0 / log(m));
        offset_.push_back(0);
}

void IndexHNSW::set_default_prob(float mL) {
    int level = 0, nn = 0;
    
    do {
        float prob = exp(-level / mL) * (1 - exp(-1 / mL));
        if (prob < 1e-9)
            break;
        assign_probs.push_back(prob);
        level += 1;
        nn += level == 0 ? m0_ : m_;
        cum_neighbor_num_.push_back(nn);
    } while (1);
    assert(level > 0);
}


int IndexHNSW::random_level() {
    double f = rand_float();
    // could be a bit faster with bissection
    for (int level = 0; level < assign_probs.size(); level++) {
        if (f < assign_probs[level]) {
            return level;
        }
        f -= assign_probs[level];
    }
    // happens with exponentially low probability
    return assign_probs.size() - 1;
}

void IndexHNSW::prepare_level() {
    // first we set the level of every vector
    levels_.resize(ntotal_);
    for (size_t i = 0; i != ntotal_; i++) {
        levels_[i] = random_level(); // calculate level of current
    }

    // now we calculate the neigbor offset of every vector
    int offset = 0;
    for (size_t i = 0; i != ntotal_; i++) {
        int pt_level = levels_[i];
        offset += cum_neighbor_num_[pt_level];   
        offset_.push_back(offset); 
        neighbors_.resize(offset, -1); // initialize as -1
    }

}

void IndexHNSW::train() {
    prepare_level();

    for (size_t idx = 0; idx < ntotal_; idx++) {
        insert(idx); // insert the i-th vector to hnsw
    }
}


void IndexHNSW::insert(size_t idx) {
    // get the vector of idx
    const float *xi = get_xb(idx);

    // if no entry point
    if (entry_point_ == -1) {
        entry_point_ = idx;
        max_level_ = levels_[idx];
        return;
    }

    
}


double IndexHNSW::rand_float() {
    return mt_() / static_cast<float>(mt_.max());
}